import numpy as np
import cv2 as cv
from matplotlib import pyplot as plt


# 通过叉积判断点在向量哪一侧(正左负右)
def cross_product(start, end, point):
    res = (end[0] - start[0]) * (point[1] - start[1]) - (end[1] - start[1]) * (point[0] - start[0])
    return res


# 图像处理的主函数
def draw_circles(ex, img):
    contours, hierarchy = cv.findContours(ex, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    # 筛选轮廓并对轮廓面积排序,第一大的为粗黑带,第二大的为细黑带
    contours = [cnt for cnt in contours if 200 < cv.contourArea(cnt) < 1000]
    area = []
    for cnt in contours:
        area.append(cv.contourArea(cnt))
    arg_area = np.argsort(area)
    thick_band = contours[arg_area[-1]]
    rect1 = cv.minAreaRect(thick_band)
    thick_center = rect1[0]
    box1 = cv.boxPoints(rect1)
    box1 = np.int64(box1)
    cv.drawContours(img, [box1], 0, (0, 0, 255), 2)
    slim_band = contours[arg_area[-2]]
    rect2 = cv.minAreaRect(slim_band)
    slim_center = rect2[0]
    box2 = cv.boxPoints(rect2)
    box2 = np.int64(box2)
    cv.drawContours(img, [box2], 0, (0, 255, 255), 2)
    plt.imshow(img)
    plt.show()
    # 确定处理区域的边框
    box = np.concatenate((box1, box2))
    rect = np.zeros((4, 2), dtype='float32')
    s = box.sum(axis=1)  # 计算x+y的值
    rect[0] = box[np.argmin(s)]  # x+y值最小的定为左上角
    rect[2] = box[np.argmax(s)]  # x+y值最大的定为右上角
    diff = np.diff(box, axis=1)  # 计算x-y
    rect[1] = box[np.argmin(diff)]  # 差值最小的点为右上角点
    rect[3] = box[np.argmax(diff)]  # 差值最大表示左下角点
    rect = sorted(rect, key=lambda x: (np.linalg.norm(x - thick_center)))
    rect = np.array(rect)
    rect = rect.reshape(2, 2, 2)
    sorted_rect = []
    for rec in rect:
        rec = sorted(rec, key=lambda x: (cross_product(thick_center, slim_center, x)))
        sorted_rect.append(rec)
    sorted_rect = np.array(sorted_rect).reshape([4, 2])  # 返回右上,左上,右下,左下点坐标

    tr, tl, br, bl = sorted_rect
    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))  # 计算底边的宽度
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))  # 计算顶边的宽度
    maxWidth = max(int(widthA), int(widthB))  # 返回最大宽度
    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))  # 计算左上角到右下角的高长度
    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))  # 计算右上角到左下角的高的长度
    maxHeight = max(int(heightA), int(heightB))  # 返回最长的高度
    dst = np.array([[maxWidth - 1, 0],  # 定义四个点，表示变换后的矩阵的角点
                    [0, 0],
                    [maxWidth - 1, maxHeight - 1],
                    [0, maxHeight - 1]], dtype='float32')
    M = cv.getPerspectiveTransform(sorted_rect, dst)  # 根据原始点和变换后的点计算透视变换矩阵M
    wrapped = cv.warpPerspective(ex, M, (maxWidth, maxHeight))  # 对原始图像，针推变换矩阵和输出图像大小进行透视变换，返回变换后的图片
    img = cv.warpPerspective(img, M, (maxWidth, maxHeight))

    contours, hierarchy = cv.findContours(wrapped, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    contours = [cnt for cnt in contours if 100 < cv.contourArea(cnt) < 1000]
    area = []
    for cnt in contours:
        area.append(cv.contourArea(cnt))
    arg_area = np.argsort(area)
    thick_band = contours[arg_area[-1]]
    rect1 = cv.minAreaRect(thick_band)
    thick_center = rect1[0]
    slim_band = contours[arg_area[-2]]
    rect2 = cv.minAreaRect(slim_band)
    slim_center = rect2[0]

    # 用霍夫变换找圆
    circles = cv.HoughCircles(wrapped, cv.HOUGH_GRADIENT, 1, 10, param1=50, param2=9, minRadius=4, maxRadius=10)
    circles = np.uint16(np.around(circles))
    circles = circles[0, :]
    assert circles.shape[0] == 6, '找到的点数量错误!'
    # 用圆到粗黑带中心点的欧几里得距离对圆进行排序&两两分组成一行
    circles = sorted(circles, key=lambda x: (np.linalg.norm(x[:2] - thick_center)))
    circles = np.array(circles)
    circles = circles.reshape(3, 2, 3)
    # 以粗黑带中心点为起点,细黑带中心点为终点确定向量,通过与圆心的叉积确定点在向量哪一侧(正左负右)来行内排序
    sorted_circles = []
    for circles_row in circles:
        circles_row = sorted(circles_row, key=lambda x: (cross_product(thick_center, slim_center, x[:2])))
        sorted_circles.append(circles_row)
    sorted_circles = np.array(sorted_circles).reshape([6, 3])
    # 判断空心还是实心
    j = 0
    for i in sorted_circles:
        # 查看圆心周围像素点值,一旦有255便判定为空心
        a = np.int8(-2)
        while 1:
            if wrapped[i[1] + a, i[0] + a] == 255:
                cv.circle(img, (i[0], i[1]), i[2], (0, 255, 0), 2)
                cv.putText(img, '0,%d' % j, (i[0], i[1] + 20), cv.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0), 1)
                j = j + 1
                break
            a = a + 1
            if a == 3:
                cv.circle(img, (i[0], i[1]), i[2], (0, 0, 255), 2)
                cv.putText(img, '1,%d' % j, (i[0], i[1] + 20), cv.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0), 1)
                j = j + 1
                break
    b, g, r = cv.split(img)
    img = cv.merge([r, g, b])
    plt.imshow(img, cmap="gray")
    plt.show()


img1 = cv.imread('ex1.jpg')
ex1 = cv.cvtColor(img1, cv.COLOR_BGR2GRAY)
img2 = cv.imread('ex2.jpg')
rows, cols = img2.shape[:2]
M = cv.getRotationMatrix2D(((cols - 1) / 2.0, (rows - 1) / 2.0), 30, 1)
img2 = cv.warpAffine(img2, M, (cols, rows))

ex2 = cv.cvtColor(img2, cv.COLOR_BGR2GRAY)
# 进行高斯模糊 核值小时会导致实心变空心
ex1 = cv.GaussianBlur(ex1, (9, 9), 0)
ex2 = cv.GaussianBlur(ex2, (9, 9), 0)
# 自适应阈值(均值) 使用高斯算法时也会导致实心变空心
ex1 = cv.adaptiveThreshold(ex1, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 15)
ex2 = cv.adaptiveThreshold(ex2, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 15)

# 调用主函数进行图像处理
draw_circles(ex1, img1)
draw_circles(ex2, img2)

# 改成rgb格式
b, g, r = cv.split(img1)
img1 = cv.merge([r, g, b])
b, g, r = cv.split(img2)
img2 = cv.merge([r, g, b])

plt.subplot(1, 2, 1), plt.imshow(img1), plt.xticks([]), plt.yticks([])
plt.subplot(1, 2, 2), plt.imshow(img2), plt.xticks([]), plt.yticks([])

plt.show()
